package com.ai.bayesi.model;


import com.ai.bayesi.core.BayesianNetwork;
import com.ai.bayesi.core.Node;

import java.util.*;

/**
 * 财富预测模型
 */
public class WealthPredictionModel {
    private BayesianNetwork network;

    public WealthPredictionModel() {
        this.network = buildWealthPredictionNetwork();
    }

    /**
     * 构建财富预测贝叶斯网络
     */
    private BayesianNetwork buildWealthPredictionNetwork() {
        BayesianNetwork bn = new BayesianNetwork();

        // 定义节点及其可能取值
        Node familyBackground = new Node("FamilyBackground", Arrays.asList("Low", "Medium", "High"));
        Node education = new Node("Education", Arrays.asList("HighSchool", "Bachelor", "Master", "PhD"));
        Node workExperience = new Node("WorkExperience", Arrays.asList("Entry", "Mid", "Senior", "Expert"));
        Node skills = new Node("Skills", Arrays.asList("Basic", "Intermediate", "Advanced"));
        Node career = new Node("Career", Arrays.asList("Business", "Tech", "Healthcare", "Education", "Arts"));
        Node investment = new Node("Investment", Arrays.asList("None", "Moderate", "Aggressive"));
        Node riskTolerance = new Node("RiskTolerance", Arrays.asList("Low", "Medium", "High"));
        Node wealth = new Node("Wealth", Arrays.asList("Poor", "Middle", "Rich"));

        // 设置节点关系
        education.addParent(familyBackground);
        skills.addParent(education);
        skills.addParent(workExperience);
        career.addParent(education);
        career.addParent(skills);
        investment.addParent(familyBackground);
        investment.addParent(career);
        riskTolerance.addParent(familyBackground);
        riskTolerance.addParent(career);
        wealth.addParent(career);
        wealth.addParent(workExperience);
        wealth.addParent(investment);
        wealth.addParent(riskTolerance);

        // 添加节点到网络
        bn.addNode(familyBackground);
        bn.addNode(education);
        bn.addNode(workExperience);
        bn.addNode(skills);
        bn.addNode(career);
        bn.addNode(investment);
        bn.addNode(riskTolerance);
        bn.addNode(wealth);

        // 设置条件概率表 (CPT)
        setFamilyBackgroundCPT(familyBackground);
        setEducationCPT(education);
        setWorkExperienceCPT(workExperience);
        setSkillsCPT(skills);
        setCareerCPT(career);
        setInvestmentCPT(investment);
        setRiskToleranceCPT(riskTolerance);
        setWealthCPT(wealth);

        return bn;
    }

    /**
     * 设置家庭背景的条件概率表
     */
    private void setFamilyBackgroundCPT(Node familyBackground) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        cpt.put(Collections.emptyList(), new double[]{0.5, 0.35, 0.15}); // P(Low)=0.5, P(Medium)=0.35, P(High)=0.15
        familyBackground.setCpt(cpt);
    }

    /**
     * 设置教育水平的条件概率表
     */
    private void setEducationCPT(Node education) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        cpt.put(Collections.singletonList("Low"), new double[]{0.6, 0.3, 0.08, 0.02});
        cpt.put(Collections.singletonList("Medium"), new double[]{0.25, 0.5, 0.2, 0.05});
        cpt.put(Collections.singletonList("High"), new double[]{0.05, 0.3, 0.45, 0.2});
        education.setCpt(cpt);
    }

    /**
     * 设置工作经验的条件概率表
     */
    private void setWorkExperienceCPT(Node workExperience) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        cpt.put(Collections.emptyList(), new double[]{0.4, 0.3, 0.2, 0.1}); // Entry, Mid, Senior, Expert
        workExperience.setCpt(cpt);
    }

    /**
     * 设置技能水平的条件概率表
     */
    private void setSkillsCPT(Node skills) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        // 填充技能水平的CPT (简化版)
        cpt.put(Arrays.asList("HighSchool", "Entry"), new double[]{0.7, 0.25, 0.05});
        cpt.put(Arrays.asList("HighSchool", "Mid"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("HighSchool", "Senior"), new double[]{0.3, 0.5, 0.2});
        cpt.put(Arrays.asList("HighSchool", "Expert"), new double[]{0.1, 0.6, 0.3});
        cpt.put(Arrays.asList("Bachelor", "Entry"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("Bachelor", "Mid"), new double[]{0.2, 0.6, 0.2});
        cpt.put(Arrays.asList("Bachelor", "Senior"), new double[]{0.1, 0.4, 0.5});
        cpt.put(Arrays.asList("Bachelor", "Expert"), new double[]{0.05, 0.25, 0.7});
        cpt.put(Arrays.asList("Master", "Entry"), new double[]{0.3, 0.5, 0.2});
        cpt.put(Arrays.asList("Master", "Mid"), new double[]{0.1, 0.6, 0.3});
        cpt.put(Arrays.asList("Master", "Senior"), new double[]{0.05, 0.3, 0.65});
        cpt.put(Arrays.asList("Master", "Expert"), new double[]{0.01, 0.15, 0.84});
        cpt.put(Arrays.asList("PhD", "Entry"), new double[]{0.2, 0.6, 0.2});
        cpt.put(Arrays.asList("PhD", "Mid"), new double[]{0.05, 0.5, 0.45});
        cpt.put(Arrays.asList("PhD", "Senior"), new double[]{0.01, 0.2, 0.79});
        cpt.put(Arrays.asList("PhD", "Expert"), new double[]{0.01, 0.1, 0.89});
        skills.setCpt(cpt);
    }

    /**
     * 设置职业类型的条件概率表
     */
    private void setCareerCPT(Node career) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        // 填充职业类型的CPT (简化版)
        cpt.put(Arrays.asList("HighSchool", "Basic"), new double[]{0.3, 0.1, 0.2, 0.3, 0.1});
        cpt.put(Arrays.asList("HighSchool", "Intermediate"), new double[]{0.25, 0.15, 0.2, 0.25, 0.15});
        cpt.put(Arrays.asList("HighSchool", "Advanced"), new double[]{0.2, 0.2, 0.2, 0.2, 0.2});
        cpt.put(Arrays.asList("Bachelor", "Basic"), new double[]{0.25, 0.2, 0.2, 0.25, 0.1});
        cpt.put(Arrays.asList("Bachelor", "Intermediate"), new double[]{0.2, 0.3, 0.2, 0.2, 0.1});
        cpt.put(Arrays.asList("Bachelor", "Advanced"), new double[]{0.15, 0.4, 0.15, 0.15, 0.15});
        cpt.put(Arrays.asList("Master", "Basic"), new double[]{0.2, 0.25, 0.25, 0.2, 0.1});
        cpt.put(Arrays.asList("Master", "Intermediate"), new double[]{0.15, 0.35, 0.25, 0.15, 0.1});
        cpt.put(Arrays.asList("Master", "Advanced"), new double[]{0.1, 0.45, 0.25, 0.1, 0.1});
        cpt.put(Arrays.asList("PhD", "Basic"), new double[]{0.15, 0.2, 0.3, 0.2, 0.15});
        cpt.put(Arrays.asList("PhD", "Intermediate"), new double[]{0.1, 0.2, 0.4, 0.15, 0.15});
        cpt.put(Arrays.asList("PhD", "Advanced"), new double[]{0.05, 0.15, 0.5, 0.15, 0.15});
        career.setCpt(cpt);
    }

    /**
     * 设置投资策略的条件概率表
     */
    private void setInvestmentCPT(Node investment) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        // 填充投资策略的CPT (简化版)
        cpt.put(Arrays.asList("Low", "Business"), new double[]{0.5, 0.3, 0.2});
        cpt.put(Arrays.asList("Low", "Tech"), new double[]{0.6, 0.3, 0.1});
        cpt.put(Arrays.asList("Low", "Healthcare"), new double[]{0.65, 0.25, 0.1});
        cpt.put(Arrays.asList("Low", "Education"), new double[]{0.7, 0.25, 0.05});
        cpt.put(Arrays.asList("Low", "Arts"), new double[]{0.8, 0.15, 0.05});
        cpt.put(Arrays.asList("Medium", "Business"), new double[]{0.3, 0.5, 0.2});
        cpt.put(Arrays.asList("Medium", "Tech"), new double[]{0.35, 0.5, 0.15});
        cpt.put(Arrays.asList("Medium", "Healthcare"), new double[]{0.4, 0.45, 0.15});
        cpt.put(Arrays.asList("Medium", "Education"), new double[]{0.45, 0.45, 0.1});
        cpt.put(Arrays.asList("Medium", "Arts"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("High", "Business"), new double[]{0.1, 0.4, 0.5});
        cpt.put(Arrays.asList("High", "Tech"), new double[]{0.15, 0.4, 0.45});
        cpt.put(Arrays.asList("High", "Healthcare"), new double[]{0.2, 0.4, 0.4});
        cpt.put(Arrays.asList("High", "Education"), new double[]{0.25, 0.45, 0.3});
        cpt.put(Arrays.asList("High", "Arts"), new double[]{0.3, 0.45, 0.25});
        investment.setCpt(cpt);
    }

    /**
     * 设置风险承受能力的条件概率表
     */
    private void setRiskToleranceCPT(Node riskTolerance) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        // 填充风险承受能力的CPT (简化版)
        cpt.put(Arrays.asList("Low", "Business"), new double[]{0.4, 0.4, 0.2});
        cpt.put(Arrays.asList("Low", "Tech"), new double[]{0.5, 0.35, 0.15});
        cpt.put(Arrays.asList("Low", "Healthcare"), new double[]{0.6, 0.3, 0.1});
        cpt.put(Arrays.asList("Low", "Education"), new double[]{0.7, 0.25, 0.05});
        cpt.put(Arrays.asList("Low", "Arts"), new double[]{0.8, 0.15, 0.05});
        cpt.put(Arrays.asList("Medium", "Business"), new double[]{0.3, 0.5, 0.2});
        cpt.put(Arrays.asList("Medium", "Tech"), new double[]{0.35, 0.5, 0.15});
        cpt.put(Arrays.asList("Medium", "Healthcare"), new double[]{0.4, 0.45, 0.15});
        cpt.put(Arrays.asList("Medium", "Education"), new double[]{0.45, 0.45, 0.1});
        cpt.put(Arrays.asList("Medium", "Arts"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("High", "Business"), new double[]{0.15, 0.45, 0.4});
        cpt.put(Arrays.asList("High", "Tech"), new double[]{0.2, 0.4, 0.4});
        cpt.put(Arrays.asList("High", "Healthcare"), new double[]{0.25, 0.45, 0.3});
        cpt.put(Arrays.asList("High", "Education"), new double[]{0.3, 0.5, 0.2});
        cpt.put(Arrays.asList("High", "Arts"), new double[]{0.35, 0.5, 0.15});
        riskTolerance.setCpt(cpt);
    }

    /**
     * 设置财富水平的条件概率表
     */
    private void setWealthCPT(Node wealth) {
        Map<List<String>, double[]> cpt = new HashMap<>();
        // 填充财富水平的CPT (简化版)
        cpt.put(Arrays.asList("Business", "Entry", "None", "Low"), new double[]{0.6, 0.35, 0.05});
        cpt.put(Arrays.asList("Business", "Entry", "None", "Medium"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("Business", "Entry", "None", "High"), new double[]{0.4, 0.45, 0.15});
        cpt.put(Arrays.asList("Business", "Entry", "Moderate", "Low"), new double[]{0.5, 0.4, 0.1});
        cpt.put(Arrays.asList("Business", "Entry", "Moderate", "Medium"), new double[]{0.3, 0.55, 0.15});
        cpt.put(Arrays.asList("Business", "Entry", "Moderate", "High"), new double[]{0.2, 0.6, 0.2});
        // 省略其他组合... (为简化代码)
        cpt.put(Arrays.asList("Tech", "Expert", "Aggressive", "High"), new double[]{0.05, 0.25, 0.7});
        cpt.put(Arrays.asList("Healthcare", "Expert", "Aggressive", "High"), new double[]{0.1, 0.3, 0.6});
        cpt.put(Arrays.asList("Education", "Expert", "Aggressive", "High"), new double[]{0.15, 0.4, 0.45});
        cpt.put(Arrays.asList("Arts", "Expert", "Aggressive", "High"), new double[]{0.2, 0.5, 0.3});
        wealth.setCpt(cpt);
    }

    /**
     * 预测一个人将来的财富水平
     */
    public Map<String, Double> predictWealth(Map<String, String> evidence) {
        Map<String, Double> result = new LinkedHashMap<>();
        List<String> wealthLevels = Arrays.asList("Poor", "Middle", "Rich");

        for (String level : wealthLevels) {
            Map<String, String> query = Collections.singletonMap("Wealth", level);
            double prob = network.infer(query, evidence);
            result.put(level, prob);
        }

        return result;
    }

    public static void main(String[] args) {
        WealthPredictionModel model = new WealthPredictionModel();

        Map<String, String> evidence = new HashMap<>();
        evidence.put("FamilyBackground", "High"); // 添加缺失的证据
        evidence.put("Education", "PhD");
        evidence.put("WorkExperience", "Expert");
        evidence.put("Skills", "Advanced"); // 添加缺失的证据
        evidence.put("Career", "Tech");
        evidence.put("RiskTolerance", "High");
        evidence.put("Investment", "Aggressive");

        // 预测财富水平
        Map<String, Double> prediction = model.predictWealth(evidence);

        // 输出结果
        System.out.println("\n财富预测结果:");
        for (Map.Entry<String, Double> entry : prediction.entrySet()) {
            System.out.printf("%-8s: %.4f%n", entry.getKey(), entry.getValue());
        }

        // 找出最高概率的财富水平
        String mostLikelyLevel = Collections.max(prediction.entrySet(), 
                Map.Entry.comparingByValue()).getKey();
        System.out.printf("\n最可能的财富水平: %s%n", mostLikelyLevel);
    }
}    